package ai.rpc.core.registry;

import ai.rpc.core.config.RegistryConfig;
import ai.rpc.core.model.ServiceMetaInfo;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.cron.CronUtil;
import cn.hutool.cron.task.Task;
import cn.hutool.json.JSONUtil;
import io.etcd.jetcd.*;
import io.etcd.jetcd.kv.GetResponse;
import io.etcd.jetcd.op.Op;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.stream.Collectors;

/**
 * @Author: mengyu
 * @Date: 2024/7/26
 */
@Slf4j
public class EtcdRegistry implements Registry {
    private Client client;
    private KV kvClient;

    // rpc根节点
    private static final String ETCD_ROOT_PATH = "/rpc/";

    private final Set<String> localRegisterNodeKeySet = new HashSet<>();

    @Override
    public void init(RegistryConfig registryConfig) {
        client = Client.builder()
                .endpoints(registryConfig.getAddress())
                .connectTimeout(Duration.ofMillis(registryConfig.getTimeout()))
                .build();
        kvClient = client.getKVClient();
        // 开启心跳
        heartBeat();
    }

    @Override
    public void register(ServiceMetaInfo serviceMetaInfo) throws ExecutionException, InterruptedException {
        // 创建Lease和KV客户端
        Lease leaseClient = client.getLeaseClient();

        // 创建30秒租约
        long leaseId = leaseClient.grant(30).get().getID();

        // 设置键值对
        String regKey = ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey();
        ByteSequence key = ByteSequence.from(regKey, StandardCharsets.UTF_8);
        ByteSequence value = ByteSequence.from(JSONUtil.toJsonStr(serviceMetaInfo), StandardCharsets.UTF_8);

        // 将键值对与租约关联起来，并设置过期时间
        PutOption putOption = PutOption.builder().withLeaseId(leaseId).build();
        kvClient.put(key, value, putOption).get();
        // 心跳检测set
        localRegisterNodeKeySet.add(regKey);

    }

    @Override
    public void unRegister(ServiceMetaInfo serviceMetaInfo) throws ExecutionException, InterruptedException {
        // 删除key-value
        kvClient.delete(ByteSequence.from(ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey(), StandardCharsets.UTF_8)).get();
        // 删除心跳检测set
        localRegisterNodeKeySet.remove(ETCD_ROOT_PATH + serviceMetaInfo.getServiceNodeKey());
    }

    @Override
    public List<ServiceMetaInfo> serviceDiscovery(String serviceKey) {

        // prefix key
        String key = ETCD_ROOT_PATH + serviceKey + "/";

        try {
            // 前缀查询
            GetOption getOption = GetOption.builder().isPrefix(true).build();
            List<KeyValue> kvs = kvClient.get(
                    ByteSequence.from(key, StandardCharsets.UTF_8),
                            getOption
            )
                    .get()
                    .getKvs();
            return kvs.stream().map(kv->{
                String value = kv.getValue().toString(StandardCharsets.UTF_8);
                return JSONUtil.toBean(value, ServiceMetaInfo.class);
            }).collect(Collectors.toList());
        } catch (Exception e) {
            throw new RuntimeException("获取服务列表失败", e);
        }
    }

    @Override
    public void destroy() {
        log.warn("当前节点下线");
        // 下线节点
        for (String s : localRegisterNodeKeySet) {
            try {
                kvClient.delete(ByteSequence.from(s,StandardCharsets.UTF_8)).get();
            }catch (Exception e){
                throw new RuntimeException(s + "节点下线失败");
            }
        }

        // 释放资源
        if(kvClient != null)
            kvClient.close();
        if(client != null)
            client.close();
    }

    @Override
    public void heartBeat() {
        CronUtil.schedule("*/10 * * * * *", new Task() {
            @Override
            public void execute() {
                for (String key : localRegisterNodeKeySet) {
                    try {
                        List<KeyValue> kvs = kvClient.get(
                                        ByteSequence.from(key, StandardCharsets.UTF_8)
                                )
                                .get()
                                .getKvs();
                        // 判断该key是否过期
                        if(CollUtil.isEmpty(kvs))
                            continue;
                        // 续期
                        KeyValue keyValue = kvs.get(0);
                        String value = keyValue.getValue().toString(StandardCharsets.UTF_8);
                        ServiceMetaInfo serviceMetaInfo = JSONUtil.toBean(value, ServiceMetaInfo.class);
                        register(serviceMetaInfo);
                        log.debug("续期成功 : {}",serviceMetaInfo.getServiceName());
                    }catch (Exception e){
                        throw new RuntimeException(key+"续签失败",e);
                    }
                }
            }
        });
        CronUtil.setMatchSecond(true);
        CronUtil.start();
    }
}
